import os
import json

# 映射类别名为ID（你可以修改）
LABEL_MAP = {
    "blue": 0,
    "red": 0,
    # 可以继续添加其他类别
}

# 必须包含的关键点标签
KEYPOINT_ORDER = ["BUFF_TOP", "BUFF_LEFT", "BUFF_CENTER", "BUFF_RIGHT", "R_CENTER"]

def normalize(value, max_value):
    return round(value / max_value, 6)

def convert_annotation(json_path, output_dir):
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    img_width = data['imageWidth']
    img_height = data['imageHeight']
    shapes = data['shapes']

    # 收集框和关键点
    rectangles = []
    keypoints = {}

    for shape in shapes:
        label = shape['label']
        points = shape['points']
        if shape['shape_type'] == 'rectangle':
            rectangles.append((label, points))
        elif shape['shape_type'] == 'point':
            if label in KEYPOINT_ORDER:
                keypoints[label] = points[0]

    if not rectangles or len(keypoints) < len(KEYPOINT_ORDER):
        print(f"[Warning] Skipping {json_path}, missing rectangle or some keypoints.")
        return

    output_lines = []
    for label, (pt1, pt2) in rectangles:
        if label not in LABEL_MAP:
            print(f"[Warning] Skipping unknown label '{label}' in {json_path}")
            continue

        class_id = LABEL_MAP[label]
        x1, y1 = pt1
        x2, y2 = pt2
        cx = normalize((x1 + x2) / 2, img_width)
        cy = normalize((y1 + y2) / 2, img_height)
        w = normalize(abs(x2 - x1), img_width)
        h = normalize(abs(y2 - y1), img_height)

        line = [str(class_id), f"{cx}", f"{cy}", f"{w}", f"{h}"]
        for key in KEYPOINT_ORDER:
            kx, ky = keypoints[key]
            line.append(str(normalize(kx, img_width)))
            line.append(str(normalize(ky, img_height)))
        output_lines.append(" ".join(line))

    # 写入输出文件
    base_filename = os.path.splitext(os.path.basename(json_path))[0]
    output_path = os.path.join(output_dir, base_filename + ".txt")
    with open(output_path, 'w') as f:
        f.write("\n".join(output_lines))
    print(f"[OK] {output_path}")

def batch_convert(json_folder, output_folder):
    os.makedirs(output_folder, exist_ok=True)
    files = [f for f in os.listdir(json_folder) if f.endswith('.json')]
    for file in files:
        convert_annotation(os.path.join(json_folder, file), output_folder)


batch_convert(r'D:\BaiduNetdiskDownload\buff_fifth\red', r'D:\BaiduNetdiskDownload\buff_fifth\red_yolo')
